Variable saving and loading

def save_variables(self, filepath, overwrite=True):
    """Saves all module variables to a `.variables.json` file.

    Args:
        filepath (str | pathlib.Path): `str` or `pathlib.Path` object.
            Path where to save the program. Must end in `.variables.json`.
        overwrite (bool): Whether we should overwrite any existing program
            at the target location, or instead ask the user
            via an interactive prompt.
    """
    filepath = file_utils.path_to_string(filepath)
    if not filepath.endswith(".variables.json"):
        raise ValueError(
            "The filepath should ends with '.variables.json', "
            f"received filepath={filepath}"
        )
    config = self.get_state_tree()
    config_string = orjson.dumps(config, option=orjson.OPT_INDENT_2).decode()
    if file_utils.exists(filepath) and not overwrite:
        io_utils.ask_to_proceed_with_overwrite(filepath)
    with open(filepath, "w") as f:
        f.write(config_string)

def get_state_tree(self):
    """Retrieves tree-like structure of program variables.

    This method allows retrieval of different program variables (trainable,
    non-trainable, optimizer, and metrics). The variables are returned in a
    nested dictionary format, where the keys correspond to the variable
    names and the values are the nested representations of the variables.

    Example:

    ```python
    program.compile(
        optimizer=synalinks.optimizers.RandomFewShot(),
        reward=synalinks.rewards.ExactMatch(),
    )
    program.fit(x=x_train, y=y_train)
    state_tree = program.get_state_tree()
    ```

    Returns:
        (dict): A dictionary containing the nested representations of the
            requested variables. The keys are the variable names, and the
            values are the corresponding nested dictionaries.
    """
    variables = {}
    variables["trainable_variables"] = self._create_nested_dict(
        self.trainable_variables
    )
    variables["non_trainable_variables"] = self._create_nested_dict(
        self.non_trainable_variables
    )
    if self.optimizer:
        variables["optimizer_trainable_variables"] = self._create_nested_dict(
            self.optimizer.trainable_variables
        )
        variables["optimizer_non_trainable_variables"] = self._create_nested_dict(
            self.optimizer.non_trainable_variables
        )
    variables["metrics_variables"] = self._create_nested_dict(self.metrics_variables)
    return variables

def load_variables(self, filepath):
    """Load all module variables from a `.variable.json` file.

    Args:
        filepath (str | pathlib.Path): `str` or `pathlib.Path` object.
            Path to load the program's variables from.
            Must end in `.variables.json`.
    """
    filepath = file_utils.path_to_string(filepath)
    if not filepath.endswith(".variables.json"):
        raise ValueError(
            "The filepath should ends with '.variables.json', "
            f"received filepath={filepath}"
        )
    with open(filepath, "rb") as f:
        state_tree_config = orjson.loads(f.read())
    self.set_state_tree(state_tree_config)

def set_state_tree(self, state_tree):
    """Assigns values to variables of the program.

    This method takes a dictionary of nested variable values, which
    represents the state tree of the program, and assigns them to the
    corresponding variables of the program. The dictionary keys represent the
    variable names (e.g., `'trainable_variables'`, `'optimizer_variables'`),
    and the values are nested dictionaries containing the variable
    paths and their corresponding values.

    The program's own variables do not depend on module names: when
    the state tree was saved from the same program built under other
    names (e.g. one program per user or conversation), they are matched
    by position, the order the program creates them in. Any of them that
    cannot be matched raises instead of being skipped.

    Args:
        state_tree (dict): A dictionary representing the state tree of the program.
            The keys are the variable names, and the values are nested
            dictionaries representing the variable paths and their values.
    """
    for k, v in state_tree.items():
        path_value_dict = self._flatten_nested_dict(v)
        if k == "trainable_variables":
            self._assign_program_variables(self.trainable_variables, v, k)
        elif k == "non_trainable_variables":
            self._assign_program_variables(self.non_trainable_variables, v, k)
        elif k == "optimizer_trainable_variables":
            if self.optimizer:
                self._assign_variable_values(
                    self.optimizer.trainable_variables, path_value_dict
                )
        elif k == "optimizer_non_trainable_variables":
            if self.optimizer:
                self._assign_variable_values(
                    self.optimizer.non_trainable_variables, path_value_dict
                )
        elif k == "metrics_variables":
            self._assign_variable_values(self.metrics_variables, path_value_dict)
        else:
            raise ValueError(f"Unknown variable name: {k}")